#李聪 2022/3/7
#使用训练好的resnet提取图像特import random
import random

import torch
import os
import torchvision
from torch.utils.data import random_split, DataLoader

from datasets import AID_Dataset


def find_same(labels,label,index):
    j=0
    for i in labels:
        if i==label and j!=index:
            return j
        j=j+1
    return -1
def find_diff(labes,label):
    n=len(labes)
    list1=[]
    for i in range(n):
        list1.append(i)
    k=0
    while True:
        num=random.choice(list1)
        if labes[num]!=label:
            return num
        k=k+1
        if k==100:
            return -1


if __name__ == '__main__':
    # 加载参数
    model_name = "model.pth"  # 参数代填
    resnet_pretrained = torch.load(model_name)
    #print(resnet_pretrained)
    resnet_pretrained.fc[6]=torch.nn.Linear(in_features=1024,out_features=48,bias=True)
    print(resnet_pretrained)
    para_list = []
    for name, module in resnet_pretrained.named_parameters():
        #print(name)
        if name == "fc.0.weight" or name == "fc.0.bias" or name == "fc.3.weight" or name == "fc.3.bias" or name == "fc.6.weight" or name == "fc.6.bias":
            # layers.append(name)
            para_list.append(module)
        else:
            module.requires_grad = False
    device = torch.device("cuda")
    root_dir = "../autodl-tmp/train"
    #root_dir = "../AID"
    label_dir_list = os.listdir(root_dir)
    AID = AID_Dataset(root_dir, label_dir_list)
    print(len(AID))
    train_dataset = AID
    print(len(train_dataset))
    #print(len(test_dataset))
    resnet_pretrained.to(device)
    loss_fn=torch.nn.TripletMarginLoss(margin=1,p=2)
    loss_fn.to(device)
    train_dataloader=DataLoader(dataset=train_dataset,batch_size=64,shuffle=True)
    #test_dataloader=DataLoader(dataset=test_dataset,batch_size=64,shuffle=True)
    epoch=50
    optimer=torch.optim.Adam(para_list,lr=0.001)
    for step in range(epoch):
        print("*******第{}轮训练开始*********".format(step))
        for data in train_dataloader:
            imgs, label = data
            imgs=imgs.to(device)
            label=label.to(device)
            optimer.zero_grad()
            outputs=resnet_pretrained(imgs)
            #print(outputs.shape)
            anchor=torch.zeros([1,48],requires_grad=True)
            anchor=anchor.to(device)
            positive=torch.zeros([1,48],requires_grad=True)
            positive=positive.to(device)
            negative=torch.zeros([1,48],requires_grad=True)
            negative=negative.to(device)
            for i in range(len(imgs)):
                pre_anchor = outputs[i]
                pre_anchor=pre_anchor.view([1,48])
                po=find_same(label,label[i],i)
                if po==-1:
                    continue
                pre_positive=outputs[po]
                pre_positive=pre_positive.view([1,48])
                ne=find_diff(label,label[i])
                if ne==-1:
                    continue
                pre_negative=outputs[ne]
                pre_negative=pre_negative.view([1,48])
                if i==0:
                    anchor=pre_anchor
                    positive=pre_positive
                    negative=pre_negative
                else:
                    anchor=torch.cat((anchor,pre_anchor))
                    positive=torch.cat((positive,pre_positive))
                    negative=torch.cat((negative,pre_negative))
            loss=loss_fn(anchor,positive,negative)
            loss.backward()
            optimer.step()
        if step>=40:
            torch.save(resnet_pretrained,"sq_model_{}".format(step))



